import os
import json
import cv2
import numpy as np


def convert_labelme_to_yolo(json_path, output_dir, class_names, image_dir):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    # 读取 JSON 内的 imagePath
    image_path = data['imagePath']

    # 确保使用完整路径
    image_path = os.path.join(image_dir, os.path.basename(image_path))

    if not os.path.exists(image_path):
        print(f"Error: 图片 {image_path} 不存在")
        return

    image = cv2.imread(image_path)
    if image is None:
        print(f"Error: 读取图片失败 {image_path}")
        return

    h, w, _ = image.shape

    # 生成 YOLO 标注文件
    yolo_txt_path = os.path.join(output_dir, os.path.basename(json_path).replace('.json', '.txt'))
    with open(yolo_txt_path, 'w') as f:
        for shape in data['shapes']:
            label = shape['label']
            points = np.array(shape['points'])
            x_min, y_min = points.min(axis=0)
            x_max, y_max = points.max(axis=0)

            x_center = (x_min + x_max) / 2 / w
            y_center = (y_min + y_max) / 2 / h
            width = (x_max - x_min) / w
            height = (y_max - y_min) / h

            class_id = class_names.index(label)
            f.write(f"{class_id} {x_center} {y_center} {width} {height}\n")


# 你的数据路径
json_dir = r"D:\0Desktop\AiModel\yolov5\dataset\labels\train"  # 这里是labelme 保存json的文件夹
image_dir = r"D:\0Desktop\AiModel\yolov5\dataset\images\train"  # 图片存放的文件夹
output_dir = r"D:\0Desktop\AiModel\yolov5\dataset\labels\train"  # 输出 YOLO 标注文件的文件夹
os.makedirs(output_dir, exist_ok=True)

# 设置类别 这里的类别要和你自己标注的一致
class_names = ["coffee"]

# 处理所有 JSON 文件
for file in os.listdir(json_dir):
    if file.endswith(".json"):
        convert_labelme_to_yolo(os.path.join(json_dir, file), output_dir, class_names, image_dir)